#!/usr/bin/env python3
import argparse
import os
import pathlib
import shutil
import subprocess
import xml.etree.ElementTree as ET

import argcomplete


def pkg_name():
    try:
        tree = ET.parse('package.xml')
    except FileNotFoundError:
        print(
            'Could not find package.xml in current working directory.\n'
            'Are you sure we are inside a ROS package?'
        )
        return None
    root = tree.getroot()
    pkg_name = root.find('name').text
    return pkg_name


def workspace_root(recursive=False):
    working_dir = os.getcwd()
    while True:
        dirs = os.listdir(working_dir)
        dirs = [
            os.path.basename(os.path.join(working_dir, d))
            for d in dirs
            if os.path.isdir(os.path.join(working_dir, d))
        ]
        if 'install' in dirs:
            print('Found install dir.')
            file_path = os.path.join(working_dir, 'install', 'setup.bash')
            if os.path.isfile(file_path):
                return working_dir
        if working_dir == '/':
            print('Could not find setup file. Giving up.')
            return None
        if not recursive:
            return None
        working_dir = os.path.dirname(working_dir)


def get_paths(root_dir, selection_arg):
    cmd = f'cd {root_dir} && colcon list --names-only {selection_arg}'
    print(cmd)
    x = subprocess.Popen(cmd, stdout=subprocess.PIPE, text=True, shell=True)
    names = []
    dirs = []
    for line in iter(x.stdout.readline, ''):
        line: str = line.replace('\n', '')
        names.append(line.removeprefix('src'))
        path = pathlib.Path(root_dir) / 'install' / line
        if path.is_dir():
            dirs.append(str(path))
        else:
            print(f'{path} does not exist')
        path = pathlib.Path(root_dir) / 'build' / line
        if path.is_dir():
            dirs.append(str(path))
        else:
            print(f'{path} does not exist')
    return dirs


def delete(paths):
    for path in paths:
        shutil.rmtree(path)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'selection',
        nargs='?',
        default='',
        choices=('all', 'this', 'up-to-this'),
    )
    argcomplete.autocomplete(parser)
    args = parser.parse_args()

    if args.selection == '' or args.selection == 'all':
        selection_arg = ''
    elif args.selection == 'this':
        selection_arg = f'--packages-select {pkg_name()}'
    elif args.selection == 'up-to-this':
        selection_arg = f'--packages-up-to {pkg_name()}'

    ws_root = workspace_root(recursive=True)
    if ws_root is None:
        print('Could not find workspace root.')
        exit(1)
    paths = get_paths(ws_root, selection_arg)
    if paths:
        print('Going to delete the following directories:')
        for path in paths:
            print(path)
        print('Continue? (y/N)')
        ans = input().lower()
        if ans == 'y':
            delete(paths)
    else:
        print('There is nothing to do!')


if __name__ == '__main__':
    main()
